package com.imyuanma.qingyun.common.client.monitor.trace;

import com.imyuanma.qingyun.common.util.StringUtil;
import com.imyuanma.qingyun.common.util.constants.CommonConstants;
import org.slf4j.MDC;


/**
 * 链路上下文
 *
 * @author wangjy
 * @date 2022/09/17 00:24:28
 */
public class TraceContext {
    /**
     * 链路上下文
     */
    private static final ThreadLocal<TraceContext> CONTEXT = new ThreadLocal<TraceContext>() {
        @Override
        protected TraceContext initialValue() {
            return new TraceContext();
        }
    };
    /**
     * 链路id
     */
    private String traceId;
    /**
     * 业务id
     */
    private String businessId;
    /**
     * 父span
     */
    private String parentSpanId;
    /**
     * 当前span
     */
    private String spanId;


    public TraceContext() {
    }

    public static TraceContext get() {
        TraceContext context = CONTEXT.get();
        if (context == null) {
            context = new TraceContext();
            CONTEXT.set(context);
        }
        return context;
    }

    public static void init(String traceId) {
        if (StringUtil.isBlank(traceId)) {
            traceId = StringUtil.getUUID();
        }
        MDC.put(CommonConstants.TRACE_KEY, traceId);
        TraceContext context = get();
        context.traceId = traceId;
    }

    public static void clear() {
        CONTEXT.remove();
        MDC.remove(CommonConstants.TRACE_KEY);
    }

    public static void traceId(String traceId) {
        get().traceId = traceId;
        MDC.put(CommonConstants.TRACE_KEY, traceId);
    }

    public static void businessId(String businessId) {
        get().businessId = businessId;
    }

    public static void parentSpanId(String parentSpanId) {
        get().parentSpanId = parentSpanId;
    }

    public static void spanId(String spanId) {
        get().spanId = spanId;
    }

    public static String getTraceId() {
        return get().traceId;
    }

    public static String getBusinessId() {
        return get().businessId;
    }

    public static String getParentSpanId() {
        return get().parentSpanId;
    }

    public static String getSpanId() {
        return get().spanId;
    }
}
